///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/wfn/coeff_obj.cpp
///
/// \brief   implementation file for the coefficient object class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<numeric>

#include<formic/wfn/coeff_obj.h>
#include<formic/wfn/configuration.h>
#include<formic/wfn/symmetry.h>
#include<formic/exception.h>
#include<formic/input/base.h>
#include<formic/numeric/numeric.h>
#include<formic/mpi/interface.h>
#include<formic/output/output.h>
#include<formic/timing/timing.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief default constructor for use by the childrens' constructors
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> formic::CoeffObj<S>::CoeffObj()
  : _name(),
    _exact_ia_list(),
    _exact_ib_list() 
{
}

template formic::CoeffObj<double>::CoeffObj();
template formic::CoeffObj<std::complex<double> >::CoeffObj();

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Copy constructor for use with the childrens' clone functions
///
/// \param[in]      other   the object to copy from
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> formic::CoeffObj<S>::CoeffObj(const CoeffObj<S> & other)
  : _name(other._name),
    _exact_ia_list(other._exact_ia_list),
    _exact_ib_list(other._exact_ib_list)
{
}

template formic::CoeffObj<double>::CoeffObj(const CoeffObj<double> & other);
template formic::CoeffObj<std::complex<double> >::CoeffObj(const CoeffObj<std::complex<double> > & other);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Throw an error explaining that a function is not implemented
///
/// \param[in]      func    the name of the function that is not implemented
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template<class S> void formic::CoeffObj<S>::not_implemented(const std::string & func) const {
  throw formic::Exception("\"%s\" function not implemented for coefficient object \"%s\"") % func % _name;
}

template void formic::CoeffObj<double>::not_implemented(const std::string &) const;
template void formic::CoeffObj<std::complex<double> >::not_implemented(const std::string &) const;

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  prepare to take an exact sample
///
/// \param[in]      userinp  the user's input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S> void formic::CoeffObj<S>::prep_for_exact_sample(const formic::InputBase & userinp) {

  // get the number of orbitals
  const int no = userinp.get<int>("nao");

  // get the number of alpha/beta electrons
  const int na = userinp.get<int>("nalpha");
  const int nb = userinp.get<int>("nbeta");

  // get a configuration object and initialize it to the user's initial configuration
  formic::Configuration config(userinp);

  // get the number of irreps
  const int nirr = userinp.get<std::vector<int> >("opirr").size();

  // get the orbital irreps
  const std::vector<int> & orb_irrep = userinp.get<std::vector<int> >("orb_irrep");

  // get the point group's direct product table
  std::vector<int> dir_prod_table;
  formic::group_direct_product_table(userinp.get<std::string>("point_group"), dir_prod_table);

  // get the irrep of the initial configuration, which sets the wavefunction's irrep
  const int target_irrep = config.get_irrep(nirr, orb_irrep, dir_prod_table);

  // get the maximum spin configuration compound index
  const int maxi = formic::pow<int>(2, no);

  // get the small coefficient cutoff
  const double coeff_cutoff = userinp.get<double>("exact_sample_cutoff");

  // First, count the non-zero configurations on the root node
  int nnonzero = 0;
  if (formic::mpi::rank() == 0) {

    // initialize a counter for the number of configurations skipped due to small coefficients
    int small_skip_count = 0;

    // vector of coefficient norms above the cutoff
    std::vector<double> coeff_norms;

    // loop over alpha configurations
    for (int ia = 0; ia < maxi; ia++) {

      // populate the alpha configuration
      config.populate_alpha_or_beta(0, ia);

      //formic::of << "    " << boost::format("%10i") % ia;
      //formic::of << "    ";
      //for(int i = 0; i < no; i++)
      //  formic::of << this->config()[2*i+0];
      //formic::of.flush();

      // skip configurations with the wrong number of alpha particles
      if ( config.nap() != na ) {
        //formic::of << "  skipped" << std::endl;
        continue;
      }

      //formic::of << std::endl;

      // loop over beta configurations
      for (int ib = 0; ib < maxi; ib++) {

        // populate the beta configuration
        config.populate_alpha_or_beta(1, ib);

        //// print configuration
        //formic::of << "    " << boost::format("%10i") % ia;
        //formic::of << "    ";
        //for(int i = 0; i < no; i++)
        //  formic::of << this->config()[2*i+0];
        //formic::of << "    " << boost::format("%10i") % ib;
        //formic::of << "    ";
        //for(int i = 0; i < no; i++)
        //  formic::of << this->config()[2*i+1];
        //formic::of.flush();

        // skip configurations with the wrong number of beta particles
        if ( config.nbp() != nb ) {
          //formic::of << "  skipped" << std::endl;
          continue;
        }

        // skip configurations with the wrong irrep
        if ( target_irrep != config.get_irrep(nirr, orb_irrep, dir_prod_table) ) {
          //formic::of << "  skipped" << std::endl;
          continue;
        }

        //// print configuration
        ////formic::of << "    " << boost::format("%10i") % ia;
        //formic::of << "    ";
        //for(int i = 0; i < no; i++)
        //  formic::of << config[2*i+0];
        ////formic::of << "    " << boost::format("%10i") % ib;
        //formic::of << "    ";
        //for(int i = 0; i < no; i++)
        //  formic::of << config[2*i+1];
        //formic::of.flush();

        formic::start_timer("exact sample prep coeff obj");

        // get the coefficient value for the configuration
        const S cv = this->get_coeff(config).get();

        //// print coefficient value
        //formic::of << "    " << formic::format_number("%20.12f", cv) << std::endl;

        formic::stop_timer("exact sample prep coeff obj");

        // skip configurations with zero coefficients
        if ( std::abs(cv) < coeff_cutoff ) {
          small_skip_count++;
          //for(int i = 0; i < no; i++)
          //  formic::of << this->config()[2*i+0];
          ////formic::of << "    " << boost::format("%10i") % ib;
          //formic::of << "    ";
          //for(int i = 0; i < no; i++)
          //  formic::of << this->config()[2*i+1];
          //formic::of << "    " << formic::format_number("%20.6e", cv) << "     skipped" << std::endl;
          continue;
        }

        // count the configuration if it's coefficient is nonzero
        nnonzero++;

        // record the coefficient norm
        coeff_norms.push_back( std::abs(cv) );

      }

    }

    // sort the coefficient norms by size (making the median easy to find)
    std::sort(coeff_norms.begin(), coeff_norms.end());

    // print some info on the sample
    formic::of << boost::format("exact sample preparation skipped %i configurations due to small coefficients") % small_skip_count << std::endl;
    formic::of << boost::format("  info on norms above cutoff:") << std::endl;
    formic::of << boost::format("         minimum = %.2e") % *coeff_norms.begin() << std::endl;
    formic::of << boost::format("         maximum = %.2e") % *coeff_norms.rbegin() << std::endl;
    formic::of << boost::format("         average = %.2e") % ( std::accumulate(coeff_norms.begin(), coeff_norms.end(), 0.0) / coeff_norms.size() ) << std::endl;
    formic::of << boost::format("          median = %.2e") % coeff_norms.at(coeff_norms.size()/2) << std::endl;
    formic::of << std::endl;

    //throw formic::Exception("Stopping after printing configurations");

  }

  // send the number of nonzero configurations to all nodes
  formic::mpi::bcast(nnonzero);

  // determine which configurations to do on this process
  const int np = formic::mpi::size();
  const int mr = formic::mpi::rank();
  const int my_first_config = ( nnonzero / np ) * mr;
  const int my_last_config = ( mr < np - 1 ? ( nnonzero / np ) * ( mr + 1 ) : nnonzero ) - 1;

  // prepare vectors to hold this process's configuration indices
  _exact_ia_list.resize(my_last_config - my_first_config + 1);
  _exact_ib_list.resize(my_last_config - my_first_config + 1);

  // create a counter for the number of configuration indices recorded
  int my_count = 0;

  // create a counter for the number of nonzero configurations seen
  int count_nonzero = 0;

  // loop over alpha configurations
  for (int ia = 0; ia < maxi; ia++) {

    // populate the alpha configuration
    config.populate_alpha_or_beta(0, ia);

    // skip configurations with the wrong number of alpha particles
    if ( config.nap() != na )
      continue;

    // loop over beta configurations
    for (int ib = 0; ib < maxi; ib++) {

      // populate the beta configuration
      config.populate_alpha_or_beta(1, ib);

      // skip configurations with the wrong number of beta particles
      if ( config.nbp() != nb )
        continue;

      // skip configurations with the wrong irrep
      if ( target_irrep != config.get_irrep(nirr, orb_irrep, dir_prod_table) )
        continue;

      // get the coefficient value for the configuration
      const S cv = this->get_coeff(config).get();

      // skip configurations with zero coefficients
      if ( std::abs(cv) < coeff_cutoff )
        continue;

      // only do this process's configurations
      if ( count_nonzero < my_first_config || count_nonzero > my_last_config ) {
        count_nonzero++;
        continue;
      }
      count_nonzero++;

      // record that this process is responsible for this configuration
      _exact_ia_list.at(my_count) = ia;
      _exact_ib_list.at(my_count) = ib;
      my_count++;

    }

  }

  // ensure the correct number of configurations were recorded
  if ( my_count != _exact_ia_list.size() )
    throw formic::Exception("Process %i recorded %i configurations in formic::VMC_Walker::prep_for_exact_sample.  It should have recorded %i.")
          % formic::mpi::rank() % my_count % _exact_ia_list.size();

}

template void formic::CoeffObj<double>::prep_for_exact_sample(const formic::InputBase &);
template void formic::CoeffObj<std::complex<double> >::prep_for_exact_sample(const formic::InputBase &);
